{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<header>用Python量化海龟交易法则</header>\n",
    "<url>https://mp.weixin.qq.com/s?__biz=MzUyMDk1MDY2MQ==&mid=2247484277&idx=1&sn=83cb004ebc7be3f6535b0433cbd81f65&chksm=f9e3c59fce944c899c3d7236c1918c612b53764cd0e6388d87efa4400bb50a5d3bf1e41eee4c&scene=178&cur_album_id=1337833439080448000#rd</url>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import baostock as bs\n",
    "import pandas as pd\n",
    "bs.login()\n",
    "\n",
    "\n",
    "#证券类型，其中1：股票，2：指数，3：其它，4：可转债，5：ETF\n",
    "rs = bs.query_stock_basic()\n",
    "data_list = []\n",
    "while (rs.error_code == '0') & rs.next():\n",
    "    # 获取一条记录，将记录合并在一起\n",
    "    data_list.append(rs.get_row_data())\n",
    "result = pd.DataFrame(data_list, columns=rs.fields)\n",
    "result.to_csv(\"./basic_info.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = result[(result['type']=='1') | (result['type']=='2') | (result['type']=='3')]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result.head(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = result[result['status']=='1']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_hist_k_data(code,start_date,end_date,frequency='d')->pd.DataFrame:\n",
    "    \"\"\"\n",
    "    获取历史K线数据\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    bs.login()\n",
    "    rs = bs.query_history_k_data_plus(code,\"date,code,open,high,low,close,preclose,volume,amount,pctChg\",start_date,end_date,frequency=frequency)\n",
    "    data_list = []\n",
    "    while (rs.error_code == '0') & rs.next():\n",
    "        # 获取一条记录，将记录合并在一起\n",
    "        data_list.append(rs.get_row_data())\n",
    "    result = pd.DataFrame(data_list, columns=rs.fields)\n",
    "    bs.logout()\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = get_hist_k_data(\"sh.000001\",\"1991-07-15\",\"2024-06-24\")\n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_data(info,type):\n",
    "    columns = \"date,code,open,high,low,close,preclose,volume,amount,pctChg\".split(',')\n",
    "    data = pd.DataFrame(columns=columns,dtype=pd.Float64Dtype)\n",
    "\n",
    "    for idx,row in info.iterrows():\n",
    "        code = row['code']\n",
    "        start = row['ipoDate']\n",
    "        end = '2024-06-24'\n",
    "        tempData = get_hist_k_data(code,start,end)\n",
    "        data = pd.concat([data,tempData],axis=0,ignore_index=True)\n",
    "\n",
    "    data.to_csv(f\"{type}.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = result[result['type'] == '2']\n",
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = result[result['type'] == '2']\n",
    "load_data(result,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from multiprocessing import Pool\n",
    "\n",
    "pool = Pool(processes=4)\n",
    "\n",
    "for i in range(1,4):\n",
    "    data = result[result['status']==1]\n",
    "    pool.apply_async(func=get_data, args=(data,i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = pd.DataFrame()\n",
    "\n",
    "for idx,row in result.iterrows():\n",
    "    code = row['code']\n",
    "    start = row['ipoDate']\n",
    "    print(start)\n",
    "    end = '2024-06-24'\n",
    "    tempData = get_hist_k_data(code,start,end)\n",
    "    data = pd.merge([data,tempData])\n",
    "\n",
    "data.head(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data.to_csv('hist_data.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#先引入后面可能用到的包（package）\n",
    "import pandas as pd  \n",
    "import numpy as np\n",
    "import talib as ta\n",
    "from datetime import datetime,timedelta\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline   \n",
    "#正常显示画图时出现的中文和负号\n",
    "from pylab import mpl\n",
    "mpl.rcParams['font.sans-serif']=['SimHei']\n",
    "mpl.rcParams['axes.unicode_minus']=False\n",
    "#使用tushare获取交易数据\n",
    "#设置token\n",
    "import baostock as bs \n",
    "#注意token更换为你在tushare网站上获取的\n",
    "\n",
    "index={'上证综指': 'sh.000001',\n",
    "        '深证成指': 'sz.399106',\n",
    "        '沪深300': 'sh.000300',\n",
    "        '创业板指': 'sz.399006',\n",
    "        '上证50': 'sh.000016',\n",
    "        '中证500': 'sh.000905',\n",
    "        '中小板指': 'sz.399005\t',\n",
    "        '上证180': 'sh.000010'}\n",
    "#获取当前交易的股票代码和名称\n",
    "def get_code():\n",
    "    df = bs.query_stock_basic(exchange='', list_status='L')\n",
    "    codes=df.ts_code.values\n",
    "    names=df.name.values\n",
    "    stock=dict(zip(names,codes))\n",
    "    #合并指数和个股成一个字典\n",
    "    stocks=dict(stock,**index)\n",
    "    return stocks    \n",
    "#获取行情数据\n",
    "def get_daily_data(stock,start,end):\n",
    "    #如果代码在字典index里，则取的是指数数据\n",
    "    code=get_code()[stock]\n",
    "    if code in index.values():\n",
    "        df=bs.query_history_k_data_plus(code,\"date,code,open,high,low,close,preclose,volume,amount,pctChg\",start, end, frequency=\"d\")\n",
    "    #否则取的是个股数据\n",
    "    else:\n",
    "        df=bs.query_history_k_data_plus(code,\"date,code,open,high,low,close,preclose,volume,amount,pctChg\",start, end, frequency=\"d\",adjustflag=\"3]2\")\n",
    "    #将交易日期设置为索引值\n",
    "    df.index=pd.to_datetime(df.trade_date)\n",
    "    df=df.sort_index()\n",
    "    #计算收益率\n",
    "    df['ret']=df.close/df.close.shift(1)-1\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hs=get_daily_data('沪深300','20180101','')[['close','open','high','low','vol']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hs=get_daily_data('沪深300','20180101','')[['close','open','high','low','vol']]\n",
    "#最近N1个交易日最高价\n",
    "hs['up']=hs.ta.MAX(hs.high,timeperiod=20).shift(1)\n",
    "#最近N2个交易日最低价\n",
    "hs['down']=ta.MIN(hs.low,timeperiod=10).shift(1)\n",
    "#每日真实波动幅度\n",
    "hs['ATR']=ta.ATR(hs.high,hs.low,hs.close,timeperiod=20)\n",
    "hs.tail()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "FAL",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
